import json
from PIL import Image
import numpy as np
import os
from config import Config


Width,Height = Config.img_width,Config.img_height
trainpath = Config.train_path
testpath = Config.test_path


def img_pretreat(img):

    array = np.asarray(img,dtype=np.float64)/127.5
    array -= 1.0
    # array = array.transpose([1,0,2])
    return array


def get_label(row_label):
    with open(Config.str_table_path, 'r', encoding='utf-8') as f:
        seq_str = f.read()
    string = row_label[0]
    label_str = [seq_str.index(item) for item in string]
    color = row_label[1]
    label_color = []
    # print(len(color))
    for item in color:
        if item == 'r':
            label_color.append(0)
        elif item == 'y':
            label_color.append(1)
        elif item == 'b':
            label_color.append(2)
        elif item == 'e':
            label_color.append(3)
        else:
            raise RuntimeError('false color')

    return np.asarray(label_str),np.asarray(label_color)


def save_file(train_split=7000):
    '''
    get_captcha中标注的图片和标签转化为npy文件保存
    :param train_split: 训练集的个数
    :return:
    '''
    inputs = []
    labels1 = []
    labels2 = []

    with open(Config.filename2labels,'r') as f:
        js = json.load(f)
    for parent,dirnames,filenames in os.walk('get_captcha'):
        for filename in filenames:
            if filename.endswith('png'):
                file = os.path.join(parent,filename)
                img = Image.open(file)
                img_data = img_pretreat(img)
                inputs.append(img_data)
                label = js[filename.split('.')[0]]
                label_str,label_color = get_label(label)
                labels1.append(label_str)
                labels2.append(label_color)
    inputs = np.array(inputs)
    labels1 = np.array(labels1,dtype=np.float)
    labels2 = np.array(labels2,dtype=np.float)
    # print(inputs[0],labels1,labels2)
    np.save(os.path.join(trainpath, 'inputs.npy'), inputs[:train_split])
    np.save(os.path.join(trainpath, 'labels_str.npy'), labels1[:train_split])
    np.save(os.path.join(trainpath, 'labels_color.npy'), labels2[:train_split])
    np.save(os.path.join(testpath, 'inputs.npy'), inputs[train_split:])
    np.save(os.path.join(testpath, 'labels_str.npy'), labels1[train_split:])
    np.save(os.path.join(testpath, 'labels_color.npy'), labels2[train_split:])


if __name__ == "__main__":
    save_file()
